import http

import requests
import uuid
import base64
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad


url = 'https://www.api-clouds.cn/apiclouds/v1/image/query'
taskId = 'F3x7bBNpEhYFBAQGYXsBB3Ja'
apiid = "154957322465723160"
apikey = "DDAA5E5497504A65B722B101C8F8CCAE"

def get_key(key):
    return key.encode('utf-8')
def encrypt(data, key):
    cipher = AES.new(get_key(key), AES.MODE_ECB)
    padded_data = pad(data.encode('utf-8'), AES.block_size)
    encrypted_data = cipher.encrypt(padded_data)
    return base64.b64encode(encrypted_data).decode('utf-8')

def generate_token(apiid, api_key):
    data = apiid + ':' + str(uuid.uuid4())
    return encrypt(data, api_key)

#获得midjourney的图片信息
def get_midjourney_task_info(taskId,url='https://www.api-clouds.cn/apiclouds/v1/image/query',apiid="154957322465723160",apikey="DDAA5E5497504A65B722B101C8F8CCAE"):
    apitoken = generate_token(apiid,apikey)
    payload = {'taskId': taskId}
    headers = {
        'Content-Type': 'application/json',
        'x-apiid': apiid,
        'x-token': apitoken
    }
    response = requests.request('POST', url, headers=headers, json=payload)
    return response.text

def submit_midjourney_task(prompt,apiid="154957322465723160",apikey="DDAA5E5497504A65B722B101C8F8CCAE"):
    apitoken = generate_token(apiid,apikey)
    conn = http.client.HTTPSConnection("www.api-clouds.cn")
    #payload自定义prompt

    #更改payload，自定义prompt
    payload = '''{
    "action": "generate",
    "prompt": "'''+prompt+'''",
    "timeout": 200,
    "callbackUrl": "https://www.fks-alibaba.com/test/v1/callback"
    }'''

    headers = {
        "Content-Type": "application/json",
        "x-token": apitoken,
        "x-apiid": apiid
    }
    conn.request("POST", "/apiclouds/v5/midjourney/imagine", payload, headers)
    res = conn.getresponse()
    data = res.read()
    print(data.decode("utf-8"))
    return data.decode("utf-8")

